import neurokit2 as nk
import matplotlib.pyplot as plt
import wfdb

# 模拟或加载真实ECG信号
record = wfdb.rdrecord(f'../data/mit-bih-arrhythmia-database-1.0.0/100', sampfrom=0, sampto=1000,
                       physical=True,
                       channels=[0])
ecg_signal = record.p_signal.T[0]

# 1. 加载示例数据
# ecg_signal = nk.ecg_simulate(duration=10, sampling_rate=360)

# 2. 预处理并定位特征点
cleaned = nk.ecg_clean(ecg_signal, sampling_rate=360)
signals, info = nk.ecg_peaks(cleaned, sampling_rate=360, method="neurokit")
# nk.ecg_plot(signals, sampling_rate=360)
# 3. 可视化R峰检测
nk.events_plot(info["ECG_R_Peaks"], cleaned)

# 4. 提取更全面的特征（如P波、QRS复合波、T波）
ecg_events = nk.ecg_findpeaks(cleaned, sampling_rate=360)
waves = nk.ecg_delineate(cleaned, ecg_events, sampling_rate=360, method="peak")

# 5. 绘制详细的心电图波形
nk.ecg_plot(waves, sampling_rate=360)
plt.show()

# 6. 计算心率变异性（HRV）指标
hrv_features = nk.hrv(ecg_events, sampling_rate=1000, show=True)
print(hrv_features)





# sampling_rate = 360
# # 方法1：使用不同的R波检测算法
# try:
#     # 尝试Pantompkins算法（对噪声更鲁棒）
#     signals, info = nk.ecg_process(ecg_signal, sampling_rate=sampling_rate, method="pantompkins1985")
#     print("Pantompkins算法处理成功!")
# except Exception as e1:
#     print(f"Pantompkins算法失败: {e1}")
#
#     # 方法2：使用Hamilton算法
#     try:
#         signals, info = nk.ecg_process(ecg_signal, sampling_rate=sampling_rate, method="hamilton2002")
#         print("Hamilton算法处理成功!")
#     except Exception as e2:
#         print(f"Hamilton算法失败: {e2}")
#
# # 方法3：手动调整R波检测灵敏度
# rpeaks = nk.ecg_findpeaks(ecg_signal, sampling_rate=sampling_rate, method="neurokit", show=True)
# print(f"检测到的R波数量: {len(rpeaks)}")
#
# # 可视化结果
# nk.ecg_plot(signals, info)
# plt.show()
#
# # 查看检测到的R波位置
# r_peaks = info['ECG_R_Peaks']
# print(f"检测到 {len(r_peaks)} 个R波")
